
#include "formant.h"

void Formant::Formant_create(double _physicalDuration, int _nFrames, double _dt, double _t1) {
    physicalDuration = _physicalDuration;
    nFrames = _nFrames;
    dt = _dt;
    t1 =_t1;
    frames.resize(_nFrames);
}
double Formant::Sampled_indexToX(int index){
    return t1 + dt * index;
}
int Formant::Sampled_xToLowIndex(double time) {
    return floor((time - t1)/dt);
}

// time 时间/s NF=1/2/3... 第1/2/3...共振峰 interpolate 是否插值
double Formant::getFormantAtX(double time, int NF, bool interpolate = true) {
    const double undefined = 0;

    if(time < 0 || time > physicalDuration) return undefined;
    if(interpolate) {
        const double index_real = (time - t1) / dt;
        const int leftIndex = floor(index_real);
        int nearIndex, farIndex;
        double phase = index_real - leftIndex;
        if(phase < 0.5) {
            nearIndex = leftIndex;
            farIndex = leftIndex + 1;
        } else {
            farIndex = leftIndex;
            nearIndex = leftIndex + 1;
            phase = 1.0 - phase;
        }
        if(nearIndex < 0 || nearIndex >= nFrames)
            return undefined;

        if(frames[nearIndex].nFormant < NF)
            return undefined;
        const double nearValue = frames[nearIndex].frequnecy[NF-1];
        if(farIndex < 0 || farIndex >= nFrames)
            return nearValue;
        if(frames[farIndex].nFormant < NF)
            return undefined;
        const double farValue = frames[farIndex].frequnecy[NF-1];
        return nearValue + phase * (farValue - nearValue);
    }

    const double index_real = (time - t1) / dt;
    const int leftIndex = floor(index_real);
    int nearIndex;
    double phase = index_real - leftIndex;
    if(phase < 0.5) {
        nearIndex = leftIndex;
    } else {
        nearIndex = leftIndex + 1;
    }
    if(nearIndex < 0 || nearIndex >= nFrames)
        return undefined;
    if(frames[nearIndex].nFormant < NF)
        return undefined;
    const double nearValue = frames[nearIndex].frequnecy[NF-1];
    return nearValue;
}

void Formant::Free(){
    for(size_t i = 0; i < frames.size(); ++ i) {
        freeVector(frames[i].frequnecy);
        freeVector(frames[i].bandwidth);
    }
    freeVector(frames);
}

/*
 * samples 采样点
 * p LPC阶数
 * nyquistFrequency
 * safetyMargin
 * */
static void burg(const vector<double> &samples, int p, double nyquistFrequency, double safetyMargin, int &nFormant,
                 vector<double> &formant, vector<double> &bandwidth, vector<complex<double>> &preRoots) {
//    cerr << "burg : start" << endl;

    const double PI = acos(-1.0);

    vector<double> coefficients;
    lpc(samples, coefficients, p);
    vector<double> polynomail;
    polynomail.push_back(1);
    for(size_t i = 0; i < coefficients.size(); ++ i) {
        polynomail.push_back(-coefficients[i]);
    }
    vector<complex<double>> rts = roots(polynomail,p, preRoots);
    Roots_fixIntoUnitCircle(rts, p);
    preRoots = rts;

//    cerr << "before nFormant" << endl;
    nFormant = 0;
//    cerr << "after nFormant" << endl;

    for(size_t i = 0; i < rts.size(); ++ i) {
        if(rts[i].imag() > 0.0) {
            double f = fabs (atan2 (rts [i].imag(), rts [i].real())) * nyquistFrequency / PI;
            if(f >= safetyMargin && f <= nyquistFrequency - safetyMargin)
                ++ nFormant;
        }
    }
    formant.resize(nFormant, 0);
    bandwidth.resize(nFormant, 0);
    int iformant = 0;
    for(size_t i = 0; i < rts.size(); ++ i) {
        if(rts[i].imag() >= 0.0) {
            double f = fabs (atan2 (rts [i].imag(), rts [i].real())) * nyquistFrequency / PI;
            if(f >= safetyMargin && f <= nyquistFrequency - safetyMargin) {
                formant[iformant]  = f;
                bandwidth[iformant]  = -log(norm(rts[i])) * nyquistFrequency / PI;
                ++ iformant;
            }
        }
    }
    assert(iformant == nFormant);
    Formant::freeVector(rts);
    Formant::freeVector(polynomail);
    Formant::freeVector(coefficients);
}

void Sound_preEmphasis (sound &data, double preEmphasisFrequency) {
    const double PI = acos(-1.0);
    const double preEmphasis = exp (-2.0 * PI * preEmphasisFrequency / data.getfrequency());
    for (int i = data.getlength()-1; i >= 1; -- i)
        data[i] -= preEmphasis * data [i - 1];
}

void Formant_sort (Formant &me) {
    for (size_t iframe = 0; iframe < me.frames.size(); iframe ++) {
        vector<double> &frameF = me.frames[iframe].frequnecy;
        vector<double> &frameB = me.frames[iframe].bandwidth;
        int n = me.frames[iframe].nFormant;
        for (int i = 0; i < n-1; i ++) {
            double min = frameF[i];
            int imin = i;
            for (int j = i + 1; j < n; j ++)
                if (frameF[j] < min) {
                    min = frameF[j];
                    imin = j;
                }
            if (imin != i) {
                const double min_bandwidth = frameB[imin];
                frameF[imin] = frameF [i];
                frameB[imin] = frameB [i];
                frameF[i] = min;
                frameB[i] = min_bandwidth;
            }
        }
    }
}

Formant Sound_to_Formant_any_inplace(sound data, double dt_in, int numberOfPoles,
        double halfdt_window, double preemphasisFrequency, double safetyMargin) {
    vector<complex<double>> preRoots = {};

    const double dt = ( dt_in > 0.0 ? dt_in : halfdt_window / 4.0 );
    const double Fs = data.getfrequency();
    const double dx = 1.0/Fs;
    const int nx = data.getlength();
    const double physicalDuration = 1.0 * nx / Fs;
    double dt_window = 2.0 * halfdt_window;
    int nFrames = 1 + floor ((physicalDuration - dt_window) / dt);
    int nsamp_window = floor (dt_window * Fs), halfnsamp_window = nsamp_window / 2;
    if(nsamp_window < numberOfPoles + 1)
        std::cerr << "Window too short.\n", assert(0);
    double t1 = 0.0 + 0.5 * (physicalDuration - 1.0/Fs - (nFrames - 1) * dt);   // centre of first frame
    if(nFrames < 1) {
        nFrames = 1;
        t1 = 0.0 + 0.5 * physicalDuration;
        dt_window = physicalDuration;
        nsamp_window = nx;
    }

    Formant thee;
    thee.Formant_create(physicalDuration, nFrames, dt, t1);
    cerr << "Formant Creat: done" << endl;

    /* Pre-emphasis. */
    Sound_preEmphasis(data,preemphasisFrequency);
    cerr << "preEmphasis: done" << endl;

    /* Gaussian window. */
    vector<double> window(nsamp_window);
    for (int i = 1; i <= nsamp_window; i ++) {
        const double imid = 0.5 * (nsamp_window + 1), edge = exp (-12.0);
        window [i-1] = (exp (-48.0 * (i - imid) * (i - imid) / (nsamp_window + 1) / (nsamp_window + 1)) - edge) / (1.0 - edge);
    }
    cerr << "Gaussian window: done \n";

    int maximumFrameLength = nsamp_window;
    vector<double> frameBuffer(maximumFrameLength);
    for (int iframe = 0; iframe < nFrames; iframe ++) {
        const double t = thee.Sampled_indexToX (iframe);
        const int leftSample = data.Sampled_xToLowIndex (t);
        const int rightSample = leftSample + 1;
        int startSample = rightSample - halfnsamp_window;
        int endSample = leftSample + halfnsamp_window;
        double maximumIntensity = 0.0;
        if(startSample < 0) startSample = 0;
        if(endSample >= nx) endSample = nx-1;

        for (int i = startSample; i <= endSample; i ++) {
            const double value = data[i];
            if (value * value > maximumIntensity)
                maximumIntensity = value * value;
        }
        if (maximumIntensity == 0.0)
            continue;   // Burg cannot stand all zeroes
        /* Copy a pre-emphasized window to a frame. */
        const int actualFrameLength = endSample - startSample + 1;   // should rarely be less than nsamp_window
        vector<double> frame(actualFrameLength);
        const int offset = startSample;
        for (int isamp = 0; isamp < actualFrameLength; isamp ++)
            frame [isamp] = data[offset + isamp] * window [isamp];

        int nF = 0;
        vector<double> fq, bd;
        burg(frame, numberOfPoles, 0.5 / dx, safetyMargin, nF,
                     fq, bd, preRoots);
        thee.frames[iframe].nFormant = nF;
        thee.frames[iframe].frequnecy = fq;
        thee.frames[iframe].bandwidth = bd;
        Formant::freeVector(fq);
        Formant::freeVector(bd);
        Formant::freeVector(frame);
//        cerr << "burg : done" << endl;
    }
    Formant_sort (thee);
    cerr << "sort : done" << endl;

    Formant::freeVector(frameBuffer);
    Formant::freeVector(window);
    vector<complex<double>>().swap(preRoots);
    return thee;
}

Formant Sound_to_Formant_any (sound me, double dt, int numberOfPoles, double maximumFrequency,
                                  double halfdt_window, double preemphasisFrequency, double safetyMargin)
{
    const double nyquist = 0.5 * me.getfrequency();
    sound sp;
    if (maximumFrequency <= 0.0 || fabs (maximumFrequency / nyquist - 1) < 1.0e-12) {
        sp = me;   // will be modified
    } else {
        sp = me.resample (maximumFrequency * 2);
        cerr << "resample done" << endl;
    }

    return Sound_to_Formant_any_inplace (sp, dt, numberOfPoles, halfdt_window, preemphasisFrequency, safetyMargin);
}

Formant Sound_to_Formant_burg(sound me, double dt, double nFormants, double maximumFrequency, double halfdt_window, double preemphasisFrequency){
    return Sound_to_Formant_any (me, dt, floor(2.0 * nFormants), maximumFrequency, halfdt_window, preemphasisFrequency, 50.0);
}
